#include "graph.h"
#include "disjoint-set.h"
#include <set>

struct SingleEdge {
  SingleEdge() { }
  SingleEdge(int i, int j, int dist) : i(i), j(j), dist(dist) { }

  bool operator==(const SingleEdge &e) const{
    if (i == e.i && j == e.j) return true;
    if (i == e.j && j == e.i) return true;
    return false;
  }
  bool operator!=(const SingleEdge &e) const {
    return !(*this == e);
  }

  bool operator<(const SingleEdge &e) const {
    if (dist < e.dist) return true;
    if (dist == e.dist && *this != e) {
      return true;
    }
    return false;
  }
  int i, j;
  int dist;
};

typedef std::set<SingleEdge> EdgeSet;
typedef std::list<SingleEdge> Result;

inline char GetNodeName(int i) { return 'a' + i; }
inline int GetNodeId(char i) { return i - 'a'; }

template<typename T>
void Print(const T &v) {
  printf("------------------------------------------\n");
  for (typename T::const_iterator it = v.begin(); it != v.end(); ++it) {
    printf("e: %c, %c, %d\n", GetNodeName(it->i), GetNodeName(it->j), it->dist);
  }
}

void Kruskal(const Graph &g, Result *result) {
  DisjointSetForest disjoint_set(g.n_vertices());
  for (size_t i = 0; i < g.n_vertices(); ++i) {
    disjoint_set.MakeSet(i);
  }

  EdgeSet edges;
  int count = 0;
  for (size_t i = 0; i < g.n_vertices(); ++i) {
    for (EdgeList::const_iterator it = g.neighbors(i).begin(); 
         it != g.neighbors(i).end(); ++it) {
      edges.insert(SingleEdge(i, it->v, it->dist));
      ++count;
    }
  }
  Print(edges);

  for (EdgeSet::const_iterator it = edges.begin(); it != edges.end(); ++it) {
    if (disjoint_set.Find(it->i) != disjoint_set.Find(it->j)) {
      result->push_back(*it);
      disjoint_set.Union(it->i, it->j);
    }
  }
}

typedef std::set<int> NodeSet;

SingleEdge FindShortestEdge(const NodeSet &nodes, const Graph &g) {
  int i, j, dist = INF;
  for (NodeSet::const_iterator it = nodes.begin(); it != nodes.end(); ++it) {
    const EdgeList &edges = g.neighbors(*it);
    for (EdgeList::const_iterator eit = edges.begin(); eit != edges.end(); 
         ++eit) {
      if (nodes.find(eit->v) != nodes.end()) continue;

      if (eit->dist < dist) {
        i = *it;
        j = eit->v;
        dist = eit->dist;
      }
    }
  }
  return SingleEdge(i, j, dist);
}

void Prim(const Graph &g, Result *result) {
  NodeSet v;
  v.insert(0);
  while (v.size() < g.n_vertices()) {
    SingleEdge edge = FindShortestEdge(v, g);
    v.insert(edge.j);
    result->push_back(edge);
  }
}

inline void AddEdge(Graph *g, char i, char j, int dist) {
  g->AddEdge(GetNodeId(i), GetNodeId(j), dist);
}

int main(int argc, const char *argv[]) {
  Graph g(9);
  AddEdge(&g, 'a', 'b', 4);
  AddEdge(&g, 'a', 'h', 8);
  AddEdge(&g, 'b', 'h', 11);
  AddEdge(&g, 'b', 'c', 8);
  AddEdge(&g, 'c', 'd', 7);
  AddEdge(&g, 'c', 'f', 4);
  AddEdge(&g, 'c', 'i', 2);
  AddEdge(&g, 'd', 'e', 9);
  AddEdge(&g, 'd', 'f', 14);
  AddEdge(&g, 'e', 'f', 10);
  AddEdge(&g, 'f', 'g', 2);
  AddEdge(&g, 'g', 'i', 6);
  AddEdge(&g, 'g', 'h', 1);
  AddEdge(&g, 'h', 'i', 7);

  Result result;
  Kruskal(g, &result);
  //Prim(g, &result);
  Print(result);
  return 0;
}

